import asyncio
import json
from typing import Any, List, Optional, Union

from pydantic import Field

from src.agent.react import ReActAgent
from src.exceptions import TokenLimitExceeded
from src.llm import LLM
from src.logger import logger
from src.prompt.toolcall import NEXT_STEP_PROMPT, SYSTEM_PROMPT
from src.schema import (
    TOOL_CHOICE_TYPE,
    AgentState,
    Memory,
    Message,
    ToolCall,
    ToolChoice,
)
from src.tool import Terminate, ToolCollection

# 错误消息常量
TOOL_CALL_REQUIRED = "Tool calls required but none provided"


class ToolCallAgent(ReActAgent):
    """工具调用代理：基于 ReAct 框架的工具/函数调用处理代理。
    
    这个代理扩展了 ReAct 框架，专门用于处理工具调用场景，提供：
    1. 智能工具选择和调用机制
    2. 工具执行结果处理和错误处理
    3. 特殊工具（如终止工具）的状态管理
    4. 多媒体内容和图像处理支持
    5. 资源清理和内存管理
    
    主要功能：
    - 自动工具发现和参数解析
    - 并发工具执行支持
    - 执行结果观察和长度限制
    - Token 限制异常处理
    - 工具调用链的状态跟踪
    
    适用场景：
    - 需要调用外部 API 或工具的任务
    - 多步骤工具组合执行
    - 需要错误恢复的复杂工作流
    """

    # 代理基本信息
    name: str = Field(default="toolcall", description="代理名称标识")
    description: str = Field(
        default="一个能够执行工具调用的智能代理",
        description="代理功能描述"
    )

    # 提示词配置
    system_prompt: str = Field(
        default=SYSTEM_PROMPT, description="系统提示词，定义代理的工具调用行为"
    )
    next_step_prompt: str = Field(
        default=NEXT_STEP_PROMPT, description="下一步提示词，指导工具选择决策"
    )

    # 核心组件
    llm: Optional[LLM] = Field(
        default_factory=LLM, description="大语言模型实例，用于工具选择推理"
    )
    memory: Memory = Field(
        default_factory=Memory, description="对话记忆存储"
    )
    state: AgentState = Field(
        default=AgentState.IDLE, description="代理执行状态"
    )

    # 工具配置
    available_tools: ToolCollection = Field(
        default_factory=lambda: ToolCollection(Terminate()),
        description="可用工具集合，默认包含终止工具"
    )
    tool_choices: TOOL_CHOICE_TYPE = Field(
        default=ToolChoice.AUTO,  # type: ignore
        description="工具选择模式：AUTO（自动）、REQUIRED（必需）、NONE（禁用）"
    )
    special_tool_names: List[str] = Field(
        default_factory=lambda: [Terminate().name],
        description="特殊工具名称列表，这些工具有特殊的状态处理逻辑"
    )

    # 工具调用状态
    tool_calls: List[ToolCall] = Field(
        default_factory=list, description="当前步骤的工具调用列表"
    )
    current_base64_image: Optional[str] = Field(
        default=None, description="当前工具执行返回的 base64 图像数据（内部使用）"
    )

    # 执行控制参数
    max_steps: int = Field(default=30, description="最大执行步数，工具调用通常需要更多步骤")
    current_step: int = Field(default=0, description="当前执行步数")
    max_observe: Optional[Union[int, bool]] = Field(
        default=None, description="工具执行结果观察的最大长度限制"
    )

    async def think(self) -> bool:
        """思考阶段：分析当前状态并选择适当的工具进行下一步操作。
        
        这个方法实现了工具调用代理的思考逻辑：
        1. 构建对话上下文和提示词
        2. 调用 LLM 分析情况并选择工具
        3. 处理响应和解析工具调用
        4. 根据工具选择模式决定后续操作
        
        Returns:
            bool: True 表示需要执行工具调用，False 表示仅需返回思考结果
            
        Raises:
            TokenLimitExceeded: 当 Token 数量超出限制时
            RuntimeError: 当 LLM 返回空响应时
        """
        # 添加下一步提示词到对话上下文
        if self.next_step_prompt:
            user_msg = Message.user_message(self.next_step_prompt)
            self.messages += [user_msg]

        try:
            # 调用 LLM 获取带有工具选项的响应
            response = await self.llm.ask_tool(
                messages=self.messages,
                system_msgs=(
                    [Message.system_message(self.system_prompt)]
                    if self.system_prompt
                    else None
                ),
                tools=self.available_tools.to_params(),
                tool_choice=self.tool_choices,
            )
        except ValueError:
            # 对于参数错误，直接向上抛出
            raise
        except Exception as e:
            # 检查是否为包含 TokenLimitExceeded 的 RetryError
            if hasattr(e, "__cause__") and isinstance(e.__cause__, TokenLimitExceeded):
                token_limit_error = e.__cause__
                logger.error(
                    f"🚨 Token 限制错误（来自 RetryError）: {token_limit_error}"
                )
                # 添加错误信息到对话历史
                self.memory.add_message(
                    Message.assistant_message(
                        f"达到最大 Token 限制，无法继续执行: {str(token_limit_error)}"
                    )
                )
                # 设置代理为完成状态
                self.state = AgentState.FINISHED
                return False
            raise

        # 提取工具调用和响应内容
        self.tool_calls = tool_calls = (
            response.tool_calls if response and response.tool_calls else []
        )
        content = response.content if response and response.content else ""

        # 记录响应信息和工具选择结果
        logger.info(f"✨ {self.name} 的思考: {content}")
        logger.info(
            f"🛠️ {self.name} 选择了 {len(tool_calls) if tool_calls else 0} 个工具使用"
        )
        
        # 在终端显示代理思考过程（如果有内容）
        if content and content.strip():
            from src.console import visualizer
            visualizer.show_agent_thought(self.name, content, "analysis")
            
        # 记录工具准备和参数信息
        if tool_calls:
            logger.info(
                f"🧰 正在准备工具: {[call.function.name for call in tool_calls]}"
            )
            logger.info(f"🔧 工具参数: {tool_calls[0].function.arguments}")

        try:
            if response is None:
                raise RuntimeError("未从 LLM 获得响应")

            # 处理不同的工具选择模式
            if self.tool_choices == ToolChoice.NONE:
                # NONE 模式：禁止使用工具
                if tool_calls:
                    logger.warning(
                        f"🤔 喔，{self.name} 在不允许使用工具时尝试使用工具！"
                    )
                if content:
                    self.memory.add_message(Message.assistant_message(content))
                    return True
                return False

            # 创建并添加助手消息到对话历史
            assistant_msg = (
                Message.from_tool_calls(content=content, tool_calls=self.tool_calls)
                if self.tool_calls
                else Message.assistant_message(content)
            )
            self.memory.add_message(assistant_msg)

            # REQUIRED 模式：必须使用工具，如果没有工具调用则在 act() 中处理
            if self.tool_choices == ToolChoice.REQUIRED and not self.tool_calls:
                return True  # 将在 act() 中处理

            # AUTO 模式：如果没有工具调用但有内容，继续处理内容
            if self.tool_choices == ToolChoice.AUTO and not self.tool_calls:
                return bool(content)

            # 返回是否有工具调用需要执行
            return bool(self.tool_calls)
        except Exception as e:
            logger.error(f"🚨 哎呀！{self.name} 的思考过程遇到了问题: {e}")
            # 将错误信息添加到对话历史
            self.memory.add_message(
                Message.assistant_message(
                    f"处理过程中遇到错误: {str(e)}"
                )
            )
            return False

    async def act(self) -> str:
        """行动阶段：执行工具调用并处理执行结果。
        
        这个方法实现了工具调用代理的行动逻辑：
        1. 验证工具调用的有效性
        2. 并发执行多个工具调用
        3. 处理执行结果和错误
        4. 管理多媒体内容（如图像）
        5. 更新对话历史
        
        Returns:
            str: 所有工具执行结果的综合描述
            
        Raises:
            ValueError: 当在 REQUIRED 模式下没有工具调用时
        """
        # 检查是否有工具调用需要执行
        if not self.tool_calls:
            if self.tool_choices == ToolChoice.REQUIRED:
                raise ValueError(TOOL_CALL_REQUIRED)

            # 如果没有工具调用，返回最后一条消息的内容
            return self.messages[-1].content or "没有内容或命令需要执行"

        # 初始化结果列表，并发执行所有工具调用
        results = []
        for command in self.tool_calls:
            # 为每个工具调用重置 base64 图像数据
            self.current_base64_image = None

            # 执行单个工具调用
            result = await self.execute_tool(command)

            # 如果设置了观察限制，截取结果
            if self.max_observe:
                result = result[: self.max_observe]

            logger.info(
                f"🎯 工具 '{command.function.name}' 完成任务！结果: {result}"
            )

            # 将工具响应添加到对话记忆
            tool_msg = Message.tool_message(
                content=result,
                tool_call_id=command.id,
                name=command.function.name,
                base64_image=self.current_base64_image,
            )
            self.memory.add_message(tool_msg)
            results.append(result)

        # 返回所有工具执行结果的组合
        return "\n\n".join(results)

    async def execute_tool(self, command: ToolCall) -> str:
        """执行单个工具调用，具有健壮的错误处理机制。
        
        这个方法负责：
        1. 验证命令格式的有效性
        2. 解析工具参数和 JSON 数据
        3. 调用工具集合中的对应工具
        4. 处理特殊工具（如终止工具）的状态变化
        5. 管理多媒体内容返回
        6. 提供统一的错误处理和日志记录
        
        Args:
            command: 要执行的工具调用对象
            
        Returns:
            str: 工具执行结果的格式化描述
        """
        # 验证命令的基本格式
        if not command or not command.function or not command.function.name:
            return "错误: 无效的命令格式"

        # 检查工具是否在可用工具集合中
        name = command.function.name
        if name not in self.available_tools.tool_map:
            return f"错误: 未知的工具 '{name}'"

        try:
            # 解析工具参数（JSON 格式）
            args = json.loads(command.function.arguments or "{}")

            # 执行工具调用
            logger.info(f"🔧 激活工具: '{name}'...")
            result = await self.available_tools.execute(name=name, tool_input=args)

            # 处理特殊工具的状态变化
            await self._handle_special_tool(name=name, result=result)

            # 检查结果是否包含 base64 图像数据
            if hasattr(result, "base64_image") and result.base64_image:
                # 存储 base64 图像数据以便后续在 tool_message 中使用
                self.current_base64_image = result.base64_image

                # 格式化显示结果
                observation = (
                    f"观察到命令 `{name}` 的执行输出:\n{str(result)}"
                    if result
                    else f"命令 `{name}` 执行完成，无输出"
                )
                return observation

            # 格式化显示结果（标准情况）
            observation = (
                f"观察到命令 `{name}` 的执行输出:\n{str(result)}"
                if result
                else f"命令 `{name}` 执行完成，无输出"
            )

            return observation
        except json.JSONDecodeError:
            error_msg = f"解析 {name} 参数错误: JSON 格式无效"
            logger.error(
                f"📝 哎呀！'{name}' 的参数不正确 - JSON 格式无效，参数:{command.function.arguments}"
            )
            return f"错误: {error_msg}"
        except Exception as e:
            error_msg = f"⚠️ 工具 '{name}' 遇到了问题: {str(e)}"
            logger.exception(error_msg)
            return f"错误: {error_msg}"

    async def _handle_special_tool(self, name: str, result: Any, **kwargs):
        """处理特殊工具的执行和状态变化。
        
        某些工具（如终止工具）具有特殊的行为，需要在执行后
        改变代理的状态或执行流程。
        
        Args:
            name: 工具名称
            result: 工具执行结果
            **kwargs: 额外参数
        """
        # 检查是否为特殊工具
        if not self._is_special_tool(name):
            return

        # 判断是否应该结束执行
        if self._should_finish_execution(name=name, result=result, **kwargs):
            # 设置代理状态为完成
            logger.info(f"🏁 特殊工具 '{name}' 已完成任务！")
            self.state = AgentState.FINISHED

    @staticmethod
    def _should_finish_execution(**kwargs) -> bool:
        """Determine if tool execution should finish the agent"""
        return True

    def _is_special_tool(self, name: str) -> bool:
        """Check if tool name is in special tools list"""
        return name.lower() in [n.lower() for n in self.special_tool_names]

    async def cleanup(self):
        """清理代理使用的工具资源。
        
        这个方法负责在代理执行完成后清理所有资源，
        包括关闭数据库连接、释放内存、清理临时文件等。
        
        注意:
        - 对于支持清理的工具，会调用其 cleanup 方法
        - 清理过程中的错误不会影响整体清理流程
        """
        logger.info(f"🧹 正在为代理 '{self.name}' 清理资源...")
        for tool_name, tool_instance in self.available_tools.tool_map.items():
            if hasattr(tool_instance, "cleanup") and asyncio.iscoroutinefunction(
                tool_instance.cleanup
            ):
                try:
                    logger.debug(f"🧼 清理工具: {tool_name}")
                    await tool_instance.cleanup()
                except Exception as e:
                    logger.error(
                        f"🚨 清理工具 '{tool_name}' 时出错: {e}", exc_info=True
                    )
        logger.info(f"✨ 代理 '{self.name}' 清理完成。")

    async def run(self, request: Optional[str] = None) -> str:
        """运行代理，完成后自动清理资源。
        
        这个方法重写了基类的 run 方法，增加了自动资源清理功能。
        无论执行成功还是失败，都会在结束时清理资源。
        
        Args:
            request: 可选的初始用户请求
            
        Returns:
            str: 代理执行结果的文本描述
        """
        try:
            # 调用父类的 run 方法执行主要逻辑
            return await super().run(request)
        finally:
            # 无论执行结果如何，都进行资源清理
            await self.cleanup()
